
#include <iostream>
#include <string>
#include <stdint.h>
#include <stdio.h>

#include <string.h>
#include <stdlib.h>
#include <openssl/evp.h>
#include <openssl/sha.h>
#include <openssl/rsa.h>


#include <boost/array.hpp>
#include <boost/asio.hpp>
#include <boost/lexical_cast.hpp>

#include <google/protobuf/message_lite.h>

#include "pam_srp.hpp"

extern "C" {
#include "srp.h"
}

#include "srp_auth.pb.h"


namespace pam_srp
{
using boost::asio::ip::tcp;
using boost::asio::detail::socket_ops::host_to_network_short;
using boost::asio::detail::socket_ops::network_to_host_short;

void phex( const unsigned char *p, int nbytes )
{
   for( int i=0; i < nbytes; i++)
   {
      printf("%x", *p);
      ++p;
   }
   printf("\n");
}
   
    
struct GuardedBuffer {
   char * buff;
   GuardedBuffer( int nbytes ) : buff( new char[nbytes] ) {}
   ~GuardedBuffer() { delete [] buff; }
};
    
    
void transact( tcp::socket & socket, 
               google::protobuf::MessageLite & req,
               google::protobuf::MessageLite & reply )
{
   reply.Clear();
        
   int msg_size = req.ByteSize();
        
   if (msg_size > 65535)
      throw SRPError( "Message size exceeds maximum length" );
        
   uint16_t send_len = host_to_network_short(msg_size);
    
   GuardedBuffer send_buff( msg_size );
        
   req.SerializeToArray( send_buff.buff, msg_size );
        
   boost::array<boost::asio::const_buffer, 2> arrs = { boost::asio::buffer( &send_len, sizeof(send_len) ),
                                                       boost::asio::buffer( send_buff.buff, msg_size ) };

   boost::asio::write( socket, arrs, boost::asio::transfer_all() );
        
   uint16_t rcv_len;
        
   boost::asio::read( socket, boost::asio::buffer( &rcv_len, sizeof(rcv_len) ), boost::asio::transfer_all() );
        
   rcv_len = network_to_host_short( rcv_len );
        
   GuardedBuffer rcv_buff( rcv_len );
        
   boost::asio::read( socket, boost::asio::buffer( rcv_buff.buff, rcv_len ) );
                
   reply.ParseFromArray( rcv_buff.buff, rcv_len );
        
   req.Clear();
}

    
struct SRPGuard
{
   SRPUser * usr;

   SRPGuard( SRPUser * u ) : usr(u) {}
   ~SRPGuard() { srp_user_delete( usr ); }
};


/******************************************************************************
 * AES
 */
unsigned char * aes_oneshot_encrypt( const unsigned char * key, int key_len,
                                     const unsigned char * iv, int iv_len,
                                     const unsigned char * data, int data_len,
                                     int * out_len)
{
   int             nalloc    = 0;
   int             npartial  = 0;
   int             nfinal    = 0;
   unsigned char * encrypted = 0;
   unsigned char   key_buff[SHA256_DIGEST_LENGTH];
   unsigned char   iv_buff[SHA256_DIGEST_LENGTH];

   *out_len = 0;
   
   SHA256( key, key_len, key_buff );
   SHA256( iv, iv_len, iv_buff );

   EVP_CIPHER_CTX ctx;

   EVP_EncryptInit(&ctx, EVP_aes_256_cbc(), key_buff, iv_buff);

   nalloc = data_len + EVP_CIPHER_CTX_block_size(&ctx);

   encrypted = (unsigned char *) malloc( nalloc );

   EVP_EncryptUpdate(&ctx, encrypted, &npartial, data, data_len);

   
   EVP_EncryptFinal_ex(&ctx, encrypted+npartial, &nfinal);

   *out_len = npartial + nfinal;
   
   return encrypted;
}

/******************************************************************************
 * RSA
 */
bool verify_sig( std::string public_key_file, const unsigned char * data, int data_len, const unsigned char * sig, unsigned int sig_len )
{
   bool            result  = false;
   EVP_PKEY    *   pkey    = 0;
   EVP_MD_CTX      md_ctx;
   int             key_len = 0;
   unsigned char   der_encoding[1024];
   const unsigned char * pbuff   = der_encoding;
   FILE        *   fp      = fopen (public_key_file.c_str(), "rb");
   
   if (fp == NULL)
      throw SRPError( std::string("Failed to open public_key_file: ") + public_key_file );

   key_len = fread(der_encoding, 1, sizeof(der_encoding), fp);
         
   fclose(fp);

   pkey = d2i_PublicKey(EVP_PKEY_RSA, NULL, &pbuff, key_len);
   
   if (!pkey)
      throw SRPError( std::string("Invalid public key file: ") + public_key_file );
   
   EVP_VerifyInit   (&md_ctx, EVP_sha1());
   EVP_VerifyUpdate (&md_ctx, data, data_len);
   
   if ( EVP_VerifyFinal (&md_ctx, sig, sig_len, pkey) == 1)
      result = true;

   EVP_PKEY_free(pkey);
   
   return result;
}




bool test_connection( SRPSettings & settings )
{
   try
   {
      boost::asio::io_service   io_service;

      tcp::resolver             resolver(io_service);
      tcp::resolver::query      query(settings.server,
                                      boost::lexical_cast<std::string>(settings.port));
      tcp::resolver::iterator   endpoint_iterator = resolver.resolve(query);
      tcp::resolver::iterator   end;

      tcp::socket               socket(io_service);
        
      boost::system::error_code error = boost::asio::error::host_not_found;
        
      while (error && endpoint_iterator != end)
      {
         socket.close();
         socket.connect(*endpoint_iterator++, error);
      }
      if (error)
         throw boost::system::system_error(error);

      return true;
   }
   catch (std::exception& e)
   {}
        
   return false;
}


    
void authenticate_from_network( SRPSettings & settings,
                                const char * username, const char * password,
                                bool & is_password_expired,
                                const char * change_pass_username, 
                                const char * change_pass_password,
                                bool update_only_if_expired )
{
   is_password_expired = false;
        
   try
   {

      boost::asio::io_service   io_service;

      tcp::resolver             resolver(io_service);

      tcp::resolver::query      query(settings.server,
                                      boost::lexical_cast<std::string>(settings.port));

      tcp::resolver::iterator   endpoint_iterator = resolver.resolve(query);
      tcp::resolver::iterator   end;

      tcp::socket               socket(io_service);
        
      boost::system::error_code error = boost::asio::error::host_not_found;
        
      while (error && endpoint_iterator != end)
      {
         socket.close();
         socket.connect(*endpoint_iterator++, error);
      }
      if (error)
         throw boost::system::system_error(error);
          
      srp_auth_proto::SRPAuthMessage req;
      srp_auth_proto::SRPAuthMessage reply;

      SRPUser * usr =  srp_user_new( (SRP_HashAlgorithm) settings.hash_alg,
                                     (SRP_NGType)        settings.ng_type,
                                     username, 
                                     (const unsigned char *)password, 
                                     strlen(password), 0, 0 );
      
      SRPGuard guard(usr);
        
      const unsigned char * bytes_A = 0;
      const unsigned char * bytes_M = 0;
      int                   len_A   = 0;
      int                   len_M   = 0;
      const char *          ignore  = 0;
        
      // Auth step 1
      srp_user_start_authentication( usr, &ignore, &bytes_A, &len_A );
        
      req.set_msg_type( srp_auth_proto::AUTHENTICATION );
        
      srp_auth_proto::Authentication * auth = req.mutable_auth();
        
      auth->set_username( username );
      auth->set_a( std::string( (const char *)bytes_A, len_A ) );
        
      transact( socket, req, reply );
          
      if (reply.has_error_message())
         throw SRPError( reply.error_message() );
        
      // Auth step 2
      if (reply.msg_type() != srp_auth_proto::AUTHENTICATION || !reply.has_auth() 
          || !reply.auth().has_s() || !reply.auth().has_b() )
         throw SRPError( "Invalid authentication message" );
            
      srp_user_process_challenge( usr, 
                                  (const unsigned char *)reply.auth().s().data(), 
                                  reply.auth().s().length(), 
                                  (const unsigned char *)reply.auth().b().data(), 
                                  reply.auth().b().length(), 
                                  &bytes_M, &len_M );
        
      req.set_msg_type( srp_auth_proto::AUTHENTICATION );
        
      auth = req.mutable_auth();
      auth->set_m( std::string( (const char *)bytes_M, len_M ) );
        
      transact( socket, req, reply );
        
      if (reply.has_error_message())
         throw SRPError( reply.error_message() );
          
      // Auth step 3
      if (reply.msg_type() != srp_auth_proto::AUTHENTICATION || !reply.has_auth() 
          || !reply.auth().has_hamk() )
         throw SRPError( "Invalid authentication message" );

      if ( !reply.has_signed_hamk() )
         throw SRPError( "SRP Server message lacks required signature" );
            
      srp_user_verify_session( usr, (const unsigned char *) reply.auth().hamk().data() );
        
      if ( !srp_user_is_authenticated(usr) )
         throw AuthenticationFailed( "Invalid Username/Password" );

      if ( !verify_sig( settings.public_key_file,
                        (const unsigned char *) reply.auth().hamk().data(),
                        reply.auth().hamk().length(),
                        (const unsigned char *) reply.signed_hamk().data(),
                        reply.signed_hamk().length() ) )
      {
         throw SRPError("SRP Server failed public-key authentication");
      }

      if (reply.auth().has_password_expired())
         is_password_expired = reply.auth().password_expired();
        
      if ( change_pass_password &&
           ( !update_only_if_expired || is_password_expired) )
      {
         const unsigned char * bytes_s;
         const unsigned char * bytes_v;
         const unsigned char * v_enc;
         const unsigned char * srp_key;
         int len_s;
         int len_v;
         int len_venc;
         int srp_key_len;
         bool add_username = true;
            
         if ( !change_pass_username ) 
         {
            add_username = false;
            change_pass_username = username;
         }
            
         srp_create_salted_verification_key( (SRP_HashAlgorithm) settings.hash_alg,
                                             (SRP_NGType)        settings.ng_type,
                                             change_pass_username, 
                                             (const unsigned char *)change_pass_password, 
                                             strlen(change_pass_password), 
                                             &bytes_s, &len_s, 
                                             &bytes_v, &len_v, 0, 0 );
                        
         req.set_msg_type( srp_auth_proto::CHANGE_PASSWORD );
         srp_auth_proto::ChangePassword *cp = req.mutable_change_password();
            
         if (add_username)
            cp->set_username( change_pass_username );

         srp_key = srp_user_get_session_key(usr, &srp_key_len);
         
         v_enc = aes_oneshot_encrypt( srp_key, srp_key_len,
                                       bytes_s, len_s,
                                      bytes_v, len_v,
                                      &len_venc );
         
         cp->set_salt( std::string( (const char *)bytes_s, len_s) );
         cp->set_verification_key( std::string((const char *)v_enc, len_venc) );
            
         transact(socket, req, reply);
            
         if (reply.has_error_message())
            throw SRPError( "Failed to change password for %s: %s", username, reply.error_message().c_str() );
      }
        
   }
   catch (boost::system::system_error & e)
   {
      throw NetworkError( e.what() );
   }
}

} // end namespace pam_srp

